/*
 * Description: process monitor dll injector
 */
#include "process_monitor_injector.h"
#include <Windows.h>
#include <TlHelp32.h>
#include <thread>
#include <vector>
#include <string>
#include <functional>
#include "log.h"
#include "performance/utils/perf_utils.h"
#include "common/guest_utils.h"
#include <tchar.h>

namespace PERF {

constexpr const TCHAR *PROCESS_MONITOR_DLL_PATH = "C:\\ProgramData\\performance\\process_monitor.dll";
const std::string EXPLORER_EXE_NAME = "explorer.exe";

ProcessMonitorInjector::ProcessMonitorInjector()
{
    processNames_.push_back(EXPLORER_EXE_NAME);
}

void ProcessMonitorInjector::DetectAndInject()
{
    LOG_INFO("[PERF - DetectAndInject] - creating tInjector");
    std::thread tInjector(&ProcessMonitorInjector::DoDetectAndInject, this);
    tInjector.detach();
}

void ProcessMonitorInjector::DoDetectAndInject()
{
    injectedPid_ = 0;
    stopDetecting_.store(false);
    while (!stopDetecting_.load()) {
        std::unique_lock<std::mutex> lock(injectMutex_);
        if (stopDetecting_.load()) {
            break;
        }
        for (SIZE_T i = 0; i < processNames_.size(); i++) {
            const TCHAR *processName = processNames_[i].c_str();
            DWORD currentPid = PERF_UTILS.GetProcessIdByName(processName);
            if (currentPid == injectedPid_ || currentPid == 0) {
                continue;
            }
            if (!RemoteThreadInject(currentPid)) {
                LOG_ERR("[PERF - DoDetectAndInject] - dll injection failed.");
                continue;
            }
            LOG_INFO("[PERF - DoDetectAndInject] - dll was injected successfully, pid: {}", currentPid);
            injectedPid_ = currentPid;
        }
        lock.unlock();
        if (injectedPid_ == 0) {
            Sleep(INITIALIZE_CHECK_SECONDS * MILLIS_PER_SECOND);
            continue;
        }
        Sleep(CHECK_SECONDS * MILLIS_PER_SECOND);
    }
    LOG_INFO("[PERF - DoDetectAndInject] - End.");
}

bool ProcessMonitorInjector::RemoteThreadInject(DWORD dwPid)
{
    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, false, dwPid);
    if (!hProcess) {
        LOG_ERR("[PERF - RemoteThreadInject] - hProcess is nullptr, errCode: {}", GetLastError());
        return false;
    }
    LOG_INFO("[PERF - RemoteThreadInject] - OpenProcess succeed");

    SIZE_T nLen = sizeof(TCHAR) * (lstrlen(PROCESS_MONITOR_DLL_PATH) + 1);
    LPVOID pBuf = VirtualAllocEx(hProcess, nullptr, nLen, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!pBuf) {
        LOG_ERR("[PERF - RemoteThreadInject] - dll injecting memory allocating failed! errCode: {}", GetLastError());
        CloseHandle(hProcess);
        return false;
    }
    LOG_INFO("[PERF - RemoteThreadInject] - dll injecting memory allocating succeed, nLen: {}", nLen);

    SIZE_T dwWrite = 0;
    if (!WriteProcessMemory(hProcess, pBuf, PROCESS_MONITOR_DLL_PATH, nLen, &dwWrite) && nLen == dwWrite) {
        LOG_ERR("[PERF - RemoteThreadInject] - dll injecting memory writing failed! errCode: {}", GetLastError());
        VirtualFreeEx(hProcess, pBuf, 0, MEM_FREE);
        CloseHandle(hProcess);
        return false;
    }
    LOG_INFO("[PERF - RemoteThreadInject] - dll injecting memory writing succeed, dwWrite: {}", dwWrite);

    HANDLE hRemoteThread = CreateRemoteThread(
        hProcess, nullptr, 0,
        (LPTHREAD_START_ROUTINE)::GetProcAddress(::GetModuleHandle("kernel32.dll"), "LoadLibraryA"), pBuf, 0, 0);
    if (!hRemoteThread) {
        LOG_ERR("[PERF - RemoteThreadInject] - hRemoteThread creation failed! errCode: {}", GetLastError());
        VirtualFreeEx(hProcess, pBuf, 0, MEM_FREE);
        CloseHandle(hProcess);
        return false;
    }
    LOG_INFO("[PERF - RemoteThreadInject] - hRemoteThread creation succeed");

    DWORD dwRemoteResult = WaitForSingleObject(hRemoteThread, INFINITE);
    if (dwRemoteResult != WAIT_OBJECT_0) {
        DWORD gle = GetLastError();
        LOG_ERR("[PERF - RemoteThreadInject] - WaitForSingleObject failed, result: {}, GLE: {}", dwRemoteResult, gle);
    } else {
        DWORD dwRemoteExitCode;
        GetExitCodeThread(hRemoteThread, &dwRemoteExitCode);
        LOG_INFO("[PERF - RemoteThreadInject] - hRemoteThread was finished, exitCode: {}, remote load library failed "
                 "if exit code is 0",
                 dwRemoteExitCode);
    }
    VirtualFreeEx(hProcess, pBuf, 0, MEM_FREE);
    CloseHandle(hRemoteThread);
    CloseHandle(hProcess);
    return true;
}

bool CheckProcessName(HANDLE hProcess)
{
    if (hProcess == INVALID_HANDLE_VALUE) {
        LOG_ERR("[PERF - CheckProcessName] - hProcess is null");
        return false;
    }
    DWORD size = MAX_PATH;
    char procName[MAX_PATH] = { 0 };
    QueryFullProcessImageName(hProcess, 0, procName, &size);
    if (std::string(procName).rfind(EXPLORER_EXE_NAME) == std::string::npos) {
        LOG_ERR("[PERF - CheckProcessName] - process name is wrong!");
        return false;
    }
    return true;
}

bool GetModuleHandle(const std::string &dllPath, DWORD pid, MODULEENTRY32 &entry)
{
    bool ret = false;
    HANDLE hModuleSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid);
    if (hModuleSnapshot == INVALID_HANDLE_VALUE || !Module32First(hModuleSnapshot, &entry)) {
        LOG_ERR("[PERF - GetModuleHandle] - get snapshot failed");
        SAFE_CLOSE_HANDLE(hModuleSnapshot, INVALID_HANDLE_VALUE);
        return ret;
    }
    do {
        if (strncmp(dllPath.c_str(), entry.szExePath, dllPath.length()) == 0 ||
            strncmp(dllPath.c_str(), entry.szModule, dllPath.length()) == 0) {
            ret = true;
            break;
        }
    } while (Module32Next(hModuleSnapshot, &entry));
    SAFE_CLOSE_HANDLE(hModuleSnapshot, INVALID_HANDLE_VALUE);
    return ret;
}

bool ProcessMonitorInjector::Uninject()
{
    std::unique_lock<std::mutex> lock(injectMutex_);
    stopDetecting_.store(true);
    DWORD pid = injectedPid_;
    if (!pid) {
        LOG_WARN("[PERF - Uninject] - dll is not injected yet!");
        pid = PERF_UTILS.GetProcessIdByName(EXPLORER_EXE_NAME.c_str());
        return false;
    }
    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, false, pid);
    if (!CheckProcessName(hProcess)) {
        SAFE_CLOSE_HANDLE(hProcess, INVALID_HANDLE_VALUE);
        return false;
    }
    MODULEENTRY32 entry = { sizeof MODULEENTRY32 };
    if (!GetModuleHandle(PROCESS_MONITOR_DLL_PATH, pid, entry)) {
        LOG_ERR("[PERF - Uninject] - get dll handle failed!");
        SAFE_CLOSE_HANDLE(hProcess, INVALID_HANDLE_VALUE);
        return false;
    }
    HANDLE hRemoteThread =
        CreateRemoteThread(hProcess, nullptr, 0, (LPTHREAD_START_ROUTINE)FreeLibrary, (LPVOID)entry.modBaseAddr, 0, 0);
    WaitForSingleObject(hRemoteThread, INFINITE);
    SAFE_CLOSE_HANDLE(hRemoteThread, INVALID_HANDLE_VALUE);
    LOG_INFO("[PERF - Uninject] - done.");
    SAFE_CLOSE_HANDLE(hProcess, INVALID_HANDLE_VALUE);
    return true;
}

}